package com.xcj.springai.config;


import cn.hutool.core.lang.UUID;
//import com.xcj.springai.chatMemory.MySqlChatMemory;
import com.xcj.springai.chatMemory.MySqlChatMemory;
import com.xcj.springai.constant.ToolContentConstant;
import com.xcj.springai.dto.ConversationDto;
import com.xcj.springai.tools.G2sSrmTool;
import com.xcj.springai.tools.NBATool;
import lombok.Getter;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.client.advisor.api.Advisor;
import org.springframework.ai.chat.client.advisor.vectorstore.QuestionAnswerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.ChatOptions;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.chat.prompt.PromptTemplate;
import org.springframework.ai.embedding.EmbeddingModel;
import org.springframework.ai.rag.advisor.RetrievalAugmentationAdvisor;
import org.springframework.ai.rag.generation.augmentation.ContextualQueryAugmenter;
import org.springframework.ai.rag.retrieval.search.VectorStoreDocumentRetriever;
import org.springframework.ai.vectorstore.SearchRequest;
import org.springframework.ai.vectorstore.VectorStore;
import org.springframework.beans.BeansException;
import org.springframework.beans.factory.InitializingBean;
import org.springframework.context.ApplicationContext;
import org.springframework.context.ApplicationContextAware;
import org.springframework.context.annotation.Bean;
import org.springframework.stereotype.Component;
import reactor.core.publisher.Flux;

import java.util.ArrayList;
import java.util.List;
import java.util.Map;


@Component
@Getter
public class Agent implements InitializingBean , ApplicationContextAware {

    private static final String systemMessage = "你是AI智能助手--小周,进行问答式的话请检索相关文档，回答问题。如果你需要使用工具，提示请选择工具模式";
    private final PromptTemplate promptTemplate = new PromptTemplate("""
            上下文如下。

            ---------------------   
            {context}
            ---------------------

            遵循以下规则：
            1. 优先使用工具处理用户提出的问题。
            2. 若找不到合适的工具，给定上下文信息并且没有先前知识回答问题。
            3. 如果答案不在上下文中，以友好的语气回答你不知道。
            4. 避免使用“根据上下文...”或“提供的信息...”这样的表述。

            问题： {query}

            答案：
            """);


//private final PromptTemplate promptTemplate = new PromptTemplate("""
//            上下文如下。
//
//            ---------------------
//            {context}
//            ---------------------
//
//            遵循以下规则：
//            1、基于本地的向量库的相关文档，直接、准确地回答用户的问题。
//            2、回答应简明扼要，避免不必要的多步骤说明。
//            3、如果需要使用工具，请直接提示用户切换到工具调用模式。
//
//            问题： {query}
//
//            答案：
//            """);

    private ChatClient chatClient;
    private final ChatModel chatModel;
//    private  ChatMemory memoryStore;
    private final MySqlChatMemory mySqlMemoryStore;
    private final EmbeddingModel embeddingModel;
    private final VectorStore vectorStore;
    private ChatOptions chatOptions;
    private ApplicationContext applicationContext;

    private ChatClient ragChatClient;
    private ChatClient toolChatClient;
    private ChatClient expanderChatClient;






    // 构造函数，注入模型、记忆、嵌入模型、向量数据库
    public Agent(ChatModel chatModel,
                 // 上下文记忆的存储方式
//                 ChatMemory memoryStore,
                 MySqlChatMemory mySqlMemoryStore,
                 EmbeddingModel embeddingModel,
                 VectorStore vectorStore) {
        this.chatModel = chatModel;
//        this.memoryStore = memoryStore;
        this.mySqlMemoryStore = mySqlMemoryStore;
        this.embeddingModel = embeddingModel;
        this.vectorStore = vectorStore;
    }


    /**
     * 构建客户端
     */
    private void initChatClient() {

        // 构建聊天记录存储器 使用的是 MySqlChatMemoryRepository 这个类  然后再下面这个方法 doChatByStreamByUserId 中会自动读取的
//        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(memoryStore).build();

        //构建聊天记录
        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(mySqlMemoryStore).build();

        //向量化搜索 advisor 版本
        SearchRequest searchRequest = SearchRequest.builder()
                .topK(3)
                .similarityThreshold(0.5)
                .build();
        QuestionAnswerAdvisor advisor = QuestionAnswerAdvisor.builder(vectorStore)
                .searchRequest(searchRequest)
                .promptTemplate(this.promptTemplate)
                // 是否允许没有 context 时仍执行
                .build();

        //向量化 RetrievalAugmentationAdvisor 版本
        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
                .documentRetriever(VectorStoreDocumentRetriever.builder()
                        .topK(3)
                        .similarityThreshold(0.5)
                        .vectorStore(this.vectorStore)
                        .build())
                .queryAugmenter(ContextualQueryAugmenter.builder()
                        .promptTemplate(this.promptTemplate)
                        .allowEmptyContext(true)
                        .build())
                .build();

        chatClient = ChatClient.builder(this.chatModel)
                // 设置系统消息
                .defaultSystem(systemMessage)
                // 注册工具
                .defaultTools(this.getTools())
                // 工具名称绑定
                .defaultToolNames(getToolNames())
                .defaultAdvisors(
                        // 聊天记录
                        chatMemoryAdvisor,
                        //RAG检索增强
//                        advisor,
                        retrievalAugmentationAdvisor,
                        new SimpleLoggerAdvisor()
                )
                .build();

    }


    public Flux<ChatResponse> doChatByStreamByUserId(ConversationDto conversation) {
        // 如果会话ID为空，则生成一个新的UUID作为会话ID
        if (conversation.getConversationId() == null) {
            conversation.setConversationId(UUID.randomUUID().toString());
        }
        // 生成最终的会话ID，格式为"会话ID:用户ID"
        String finalConversationId = conversation.getConversationId() + ":" + conversation.getUserId();


        // 根据用户输入文本生成向量（用于向量检索）
//        float[] queryVector = embeddingModel.embed(conversation.getContents());

//        // 构建向量检索请求
//        SearchRequest searchRequest = SearchRequest.builder()
//                .query(conversation.getContents())
//                .topK(3)  // 限制最多返回3条(3)
//                .build();
//
//        // 执行相似度搜索，获得相关文档
//        List<Document> relatedDocs = vectorStore.similaritySearch(searchRequest);
//
          //把检索到的文档内容拼接成上下文，供系统参考
//        StringBuilder contextBuilder = new StringBuilder();
//        if (relatedDocs != null) {
//            for (Document doc : relatedDocs) {
//                String text = doc.getText();
//                if (text != null) {
//                    contextBuilder.append(text).append("\n");
//                }
//            }
//        }

        // ========== 构造 Prompt（提示信息） ==========
        // 创建系统消息
        Message systemMsg = new SystemMessage(systemMessage);
//        Message systemMsg = new SystemMessage(systemMessage + "\n相关文档：\n" + contextBuilder.toString());

        // 创建用户消息 用户输入内容，作为用户消息
        Message userMsg = new UserMessage(conversation.getContents());

        // 创建提示信息，包含系统消息和用户消息
        Prompt prompt = new Prompt(List.of(systemMsg, userMsg), chatOptions);


        // 调用聊天客户端进行流式聊天
        return chatClient
                // 构建提示请求
                .prompt()
                // 设置用户输入内容
                .user(conversation.getContents())
                // 设置 advisor，用于记忆机制中的会话ID标识
                .advisors(a -> {
                    a.param(ChatMemory.CONVERSATION_ID, finalConversationId);
//                             a.param(CHAT_MEMORY_CONVERSATION_ID_KEY, finalConversationId)
//                            .param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 10);
                })

                //传参工具参数
//                .toolContext(Map.of(
//                        ToolContentConstant.USER_ID, conversation.getUserId()
//                ))
                // 表示使用流式返回结果
                .stream()
                // 返回 Flux<ChatResponse>，用于异步响应
                .chatResponse();
    }


    /**
     * 初始化工具调用的 ChatClient
     */
//    private ChatClient initToolChatClient() {
//        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(mySqlMemoryStore).build();
//
//        return ChatClient.builder(this.chatModel)
//                .defaultSystem(
//                        "你是一个工具调用助手，仅通过系统已注册的工具来协助用户完成任务。\n" +
//                                "如果用户的问题不能通过工具处理，请温和地回复：\n" +
//                                "“这个问题我暂时无法处理，请继续提问或切换到问答模式。”\n" +
//                                "不要尝试直接回答问题内容。”。")
//                .defaultTools(this.getTools())
//                .defaultToolNames(this.getToolNames())
//                .defaultAdvisors(chatMemoryAdvisor, new SimpleLoggerAdvisor())
//                .build();
//    }

    /**
     * 初始化问答式聊天的 ChatClient
     */
//    private ChatClient initRagChatClient() {
//        MessageChatMemoryAdvisor chatMemoryAdvisor = MessageChatMemoryAdvisor.builder(mySqlMemoryStore).build();
//
//        Advisor retrievalAugmentationAdvisor = RetrievalAugmentationAdvisor.builder()
//                .documentRetriever(VectorStoreDocumentRetriever.builder()
//                        .similarityThreshold(0.5)
//                        .topK(5)
//                        .vectorStore(this.vectorStore)
//                        .build())
//                .queryAugmenter(ContextualQueryAugmenter.builder()
//                        .promptTemplate(this.promptTemplate)
//                        .allowEmptyContext(true)
//                        .build())
//                .build();
//
//        return ChatClient.builder(this.chatModel)
//                .defaultSystem(Agent.systemMessage)
//                .defaultAdvisors(chatMemoryAdvisor, retrievalAugmentationAdvisor, new SimpleLoggerAdvisor())
//                .build();
//    }

    private ChatClient initExpanderChatClient() {
        return ChatClient.builder(this.chatModel)
                .defaultSystem("你是一个问题扩展助手，根据用户输入生成多个相关的问题变体。")
                .defaultAdvisors(new SimpleLoggerAdvisor())
                .build();
    }


//    public Flux<ChatResponse> doChatByStreamByUserId(ConversationDto conversationRequest) {
//        if (conversationRequest.getConversationId() == null) {
//            conversationRequest.setConversationId(UUID.randomUUID().toString());
//        }
//
//        String content = conversationRequest.getContents();
//        ChatClient selectedClient;
//        String prefix;
//
//        if (content.startsWith("[1]")) {
//            selectedClient = this.ragChatClient;
//            prefix = "RAG";
//            content = content.substring(3);
//        } else if (content.startsWith("[2]")) {
//            selectedClient = this.toolChatClient;
//            prefix = "TOOL";
//            content = content.substring(3);
//        } else {
//            selectedClient = this.ragChatClient;
//            prefix = "RAG";
//        }
//
//        // conversationId 格式：RAG:abc123:user001
//        String finalConversationId = prefix + ":" + conversationRequest.getConversationId() + ":" + conversationRequest.getUserId();
//
//        return selectedClient
//                .prompt()
//                .user(content)
//                .advisors(a -> a.param(ChatMemory.CONVERSATION_ID, finalConversationId))
//                // 传参工具参数
//                .toolContext(Map.of("userId", conversationRequest.getUserId(),"extraParams", conversationRequest.getExtraParams()))
//                .stream()
//                .chatResponse();
//    }








    /**
     * Spring 初始化时自动调用，初始化 ChatClient
     */
    @Override
    public void afterPropertiesSet() throws Exception {
        initChatClient();
    }


    /**
     * 多个ChatClient时用
     * Spring 初始化时自动调用，初始化 ChatClient
     */
//    @Override
//    public void afterPropertiesSet() {
//        this.ragChatClient = initRagChatClient();
//        this.toolChatClient = initToolChatClient();
//        this.expanderChatClient = initExpanderChatClient();
//    }


    /**
     * 获取工具类数组（例如业务系统工具）
     */

    private Object[] getTools() {
        List<Object> tools = new ArrayList<>();
        tools.add(this.applicationContext.getBean(G2sSrmTool.class));
        tools.add(this.applicationContext.getBean(NBATool.class));
        return tools.toArray();
    }

    /**
     * 设置 Spring 上下文，用于获取 Bean
     */
    @Override
    public void setApplicationContext(ApplicationContext applicationContext) throws BeansException {
        this.applicationContext = applicationContext;
    }

    /**
     * 获取工具名称
     * 优先推荐/引导 AI 调用这些函数
     */
    private String[] getToolNames() {
        return new String[] {
                "getNBATeam"
        };
    }
}
